import torch
import torch.nn.functional as F
import SettingUp as s
import BuildingTheNetWorking as B
import Training as T
import matplotlib.pyplot as plt

def test():
  B.network.eval()
  test_loss = 0
  correct = 0
  with torch.no_grad():
    for data, target in s.test_loader:
      output = B.network(data)
      test_loss += F.nll_loss(output, target, size_average=False).item()
      pred = output.data.max(1, keepdim=True)[1]
      correct += pred.eq(target.data.view_as(pred)).sum()
  test_loss /= len(s.test_loader.dataset)
  T.test_losses.append(test_loss)
  print('\nTest set: Avg. loss: {:.4f}, Accuracy: {}/{} ({:.0f}%)\n'.format(
    test_loss, correct, len(s.test_loader.dataset),
    100. * correct / len(s.test_loader.dataset)))


test()
for epoch in range(1, s.n_epochs + 1):
  T.train(epoch)
  test()

fig = plt.figure()
plt.plot(T.train_counter, T.train_losses, color='blue')
plt.scatter(T.test_counter, T.test_losses, color='red')
plt.legend(['Train Loss', 'Test Loss'], loc='upper right')
plt.xlabel('number of training examples seen')
plt.ylabel('negative log likelihood loss')
plt.show()